from pandas import DataFrame
from pandas import Series
from pandas import concat
from pandas import read_csv
from pandas import datetime
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import LSTM
from math import sqrt
import matplotlib
from matplotlib import pyplot
import numpy
import numpy as np
from keras.callbacks import EarlyStopping
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score


# date-time parsing function for loading the dataset
def parser(x):
    return datetime.strptime('190'+x, '%Y-%m')


def create_dataset(dataset, scale):
    train_size = int(len(dataset) * scale)
    # 将数据集分为两部分
    train, test = dataset[0:train_size], dataset[train_size:]
    return train, test


# frame a sequence as a supervised learning problem
def timeseries_to_supervised(data, lag=1):
    df = DataFrame(data)
    columns = [df.shift(i) for i in range(1, lag+1)]
    columns.append(df)
    df = concat(columns, axis=1)
    df = df.drop(0)
    return df


# difference dataset
def difference(data, interval):
    return [data[i] - data[i - interval] for i in range(interval, len(data))]


# invert difference
def invert_difference(orig_data, diff_data, interval):
    return [diff_data[i-interval] + orig_data[i-interval] for i in range(interval, len(orig_data))]


# scale train and test data to [-1, 1]
def scale(train, test):
    # fit scaler
    scaler = MinMaxScaler(feature_range=(-1, 1))
    scaler = scaler.fit(train)
    # transform train
    train = train.reshape(train.shape[0], train.shape[1])
    train_scaled = scaler.transform(train)
    # transform test
    test = test.reshape(test.shape[0], test.shape[1])
    test_scaled = scaler.transform(test)
    return scaler, train_scaled, test_scaled


# transform series into train and test sets for supervised learning
def prepare_data(series, test_percent, n_lag):
    # extract raw values
    raw_values = series.values
    # transform data to be stationary
    diff_series = np.array(difference(raw_values, 1))
    diff_values = diff_series
    diff_values = diff_values.reshape(len(diff_values), 1)
    # rescale values to -1, 1
    scaler = MinMaxScaler(feature_range=(-1, 1))
    scaled_values = scaler.fit_transform(diff_values)
    scaled_values = scaled_values.reshape(len(scaled_values), 1)
    # transform into supervised learning problem X, y
    supervised = timeseries_to_supervised(scaled_values, n_lag)
    supervised_values = supervised.values
    # split into train and test sets
    train, test = create_dataset(supervised_values,test_percent)
    return scaler, train, test


# inverse scaling for a forecasted value
def invert_scale(scaler, yhat):
    return scaler.inverse_transform(yhat)


# fit an LSTM network to training data
def fit_lstm(train, n_lag, batch_size, nb_epoch, neurons):
    X, y = train[:, 0:n_lag], train[:, n_lag:]
    X = X.reshape(X.shape[0], 1, X.shape[1])
    model = Sequential()
    # model.add(LSTM(neurons, batch_input_shape=(batch_size, X.shape[1], X.shape[2]), stateful=True))
    model.add(LSTM(neurons, input_shape=(X.shape[1], X.shape[2]), stateful=False)) #false的效果更好
    model.add(Dense(1))
    model.summary()
    model.compile(loss='mean_squared_error', optimizer='adam')
    model.fit(X, y, epochs=nb_epoch, batch_size=batch_size, verbose=1, shuffle=False)

    # for i in range(nb_epoch):
    #     model.fit(X, y, epochs=1, batch_size=batch_size, verbose=1, shuffle=False)
    #     # model.reset_states()
    return model


def mean_absolute_percentage_error(y_true, y_pred):
    # 平均绝对百分比误差（MAPE）的计算
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    return np.mean(np.abs((y_true - y_pred) / y_true)) * 100


# evaluate the RMSE for each forecast time step
def evaluate_forecasts(test, forecasts):
    rmse = sqrt(mean_squared_error(test, forecasts))
    mae = mean_absolute_error(test, forecasts)
    mape = mean_absolute_percentage_error(test, forecasts)
    r2 = r2_score(test, forecasts)
    print(' RMSE: %f' % rmse)
    print(' MAE: %f' % mae)
    print(' MAPE: %f' % mape)
    print(' R2: %f' % r2)

    return rmse,mae,mape,r2


# run a repeated experiment
def experiment(repeats, series, epochs, batch_size, n_neurons, test_percent, n_lag):
    # prepare data
    scaler, train, test = prepare_data(series, test_percent, n_lag)

    # run experiment
    r2_scores = list()
    rmse_scores = list()
    mape_scores = list()
    mae_scores = list()
    for r in range(repeats):

        # fit the model
        # train_trimmed = train[2:, :]
        # print(train_trimmed)
        lstm_model = fit_lstm( train, n_lag, batch_size, epochs, n_neurons)

        # forecast the entire training dataset to build up state for forecasting
        train_reshaped = train[:, 0].reshape(len( train), 1, 1)
        lstm_model.predict(train_reshaped, batch_size=batch_size)

        # forecast test dataset
        test_reshaped = test[:,0:-1]
        test_reshaped = test_reshaped.reshape(len(test_reshaped), 1, 1)
        output = lstm_model.predict(test_reshaped, batch_size=batch_size)
        output = invert_scale(scaler,output)
        rmse,mae,mape,r2 = evaluate_forecasts(test[:, 1:2], output)

        rmse_scores.append(rmse)
        mae_scores.append(mae)
        mape_scores.append(mape)
        r2_scores.append(r2)
    return r2_scores

if __name__ == '__main__':
    # load dataset
    series = read_csv(
        'E:\lyf_ML_Drought\coding\ML_Drought_Prediction\indices_caculate\\result\ROW_SPEI\ROW_SPEI-12\SPEI-12_52533.txt',
        header=None, names=('TIME', 'SPEI-12'))
    series = series.set_index(['TIME'], drop=True)  # 把日期作为索引

    # experiment
    repeats = 3
    n_lag = 1
    test_percent = 0.9
    n_batch = 50
    n_neurons = 1
    results = DataFrame()
    # vary training epochs
    epochs = [3]
    for e in epochs:
        results[str(e)] = experiment(repeats, series, e, n_batch, n_neurons, test_percent, n_lag)
    # summarize results
    print(results.describe())

    # save boxplot
    results.boxplot()
    pyplot.savefig('boxplot/test_pochs.png')